#!/usr/bin/env python3
import asyncio
import time
from logger import Logger
import telnetlib3

class Connection:
    def __init__(self, name, ip, port, logger, user_id: str):
        self.name = name
        self.ip = ip
        self.port = port
        self.logger = logger
        self.reader = None
        self.writer = None
        self.is_connected = False
        self.connect_time = None
        self.session = None
        self.user_id = user_id

    async def connect(self):
        try:
            if self.is_connected:
                raise Exception(f'连接 {self.name} 已存在')

            # 使用asyncio.wait_for设置超时，因为telnetlib3.open_connection不直接支持timeout参数
            self.reader, self.writer = await asyncio.wait_for(
                telnetlib3.open_connection(self.ip, self.port),
                timeout=10
            )
            self.is_connected = True
            self.connect_time = time.time()
            
            # 注册日志文件
            self.logger.register_connection(self.name, self.ip, self.port)
            
            log_msg = f'成功连接到 {self.ip}:{self.port}'
            self.logger.log(self.name, log_msg)
            import sys
            print(log_msg, file=sys.stderr)
            return True
        except Exception as e:
            log_msg = f'连接失败: {str(e)}'
            self.logger.log(self.name, log_msg)
            import sys
            print(log_msg, file=sys.stderr)
            self.is_connected = False
            return False

    async def send_command(self, command):
        try:
            if not self.is_connected or self.reader is None or self.writer is None:
                raise Exception(f'连接 {self.name} 未建立')

            # 发送命令
            self.writer.write(command + '\n')
            await self.writer.drain()
            await asyncio.sleep(1)  # 等待命令执行

            # 读取输出
            output = []
            try:
                while True:
                    line = await asyncio.wait_for(self.reader.readline(), timeout=1.0)
                    if not line:
                        break
                    output.append(line)
            except asyncio.TimeoutError:
                pass  # 超时表示没有更多输出

            output_str = ''.join(output).strip()

            # 记录命令和输出
            self.logger.log(self.name, f'发送命令: {command}')
            self.logger.log(self.name, f'命令输出: {output_str}')

            return output_str
        except Exception as e:
            error_msg = f'发送命令失败: {str(e)}'
            self.logger.log(self.name, error_msg)
            raise Exception(error_msg)

    async def close(self):
        try:
            if self.is_connected and self.writer is not None:
                self.writer.close()
                # 处理wait_closed方法可能不存在的情况
                try:
                    await self.writer.wait_closed()
                except AttributeError:
                    # 如果没有wait_closed方法，等待一小段时间
                    await asyncio.sleep(0.5)
                self.is_connected = False
                log_msg = f'连接已关闭'
                self.logger.log(self.name, log_msg)
                
                # 取消注册日志文件
                self.logger.unregister_connection(self.name)
                
                return True
            return False
        except Exception as e:
            log_msg = f'关闭连接失败: {str(e)}'
            self.logger.log(self.name, log_msg)
            return False

    def is_alive(self):
        return self.is_connected

    def get_info(self):
        return {
            'name': self.name,
            'ip': self.ip,
            'port': self.port,
            'status': '活跃' if self.is_connected else '已关闭',
            'connect_time': self.connect_time
        }